#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2021/11/15 14:07
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : agglomerative_cllustering.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""


def load_data(data_path, test_size=0.2):
    """
    加载Excel数据集
    :param data_path: Excel文件路径
    :param test_size: 测试集比例大小
    :return: 训练集，测试集
    """
    import pandas as pd
    dataset = pd.read_excel(data_path, header=0, names=list(range(1, 5)), usecols=list(range(1, 5)))
    return dataset


def train(dataset):
    from joblib import dump
    from sklearn.cluster import AgglomerativeClustering
    agglomeration_cluster_model = AgglomerativeClustering(n_clusters=4)
    dump(agglomeration_cluster_model, "../models/agglomeration_cluster.model")  # 保存聚类模型对象
    return agglomeration_cluster_model.fit_predict(dataset)


def show(x_col, y_col, target_pred, title="", font=None):
    from matplotlib import font_manager
    import matplotlib.pyplot as plt
    font = font_manager.FontProperties(fname=font)
    plt.title(title, fontproperties=font)
    plt.scatter(x_col, y_col, c=target_pred, cmap="brg", label="samples")
    plt.legend()
    plt.show()


if __name__ == '__main__':
    from numpy import array

    DATA_PATH = "../dataset/chuanqishuju.xlsx"
    data = load_data(DATA_PATH)
    label_pred = train(data)
    print(label_pred)
    array = array(data)
    x, y = array[:, 2], array[:, 1]
    show(x, y, label_pred)
